package com.cyx.exercise.ollama;

/**
 * Ollama响应处理器
 * 专门处理包含思考过程的模型响应（如qwen3等）
 */
public class OllamaResponseProcessor {
    
    /**
     * 处理原始响应，提取实际答案
     */
    public static ProcessedResponse processResponse(String rawResponse) {
        if (rawResponse == null || rawResponse.trim().isEmpty()) {
            return new ProcessedResponse("", "", rawResponse);
        }
        
        // 检查是否包含思考过程
        if (rawResponse.contains("<think>") && rawResponse.contains("</think>")) {
            return extractThinkingResponse(rawResponse);
        } else {
            // 直接返回原始响应
            return new ProcessedResponse(rawResponse.trim(), "", rawResponse);
        }
    }
    
    /**
     * 提取包含思考过程的响应
     */
    private static ProcessedResponse extractThinkingResponse(String rawResponse) {
        try {
            // 找到最后一个</think>标签
            int thinkEndIndex = rawResponse.lastIndexOf("</think>");
            
            if (thinkEndIndex == -1) {
                // 没有结束标签，可能响应被截断
                return new ProcessedResponse("", rawResponse, rawResponse);
            }
            
            // 提取思考过程
            int thinkStartIndex = rawResponse.indexOf("<think>");
            String thinkingProcess = "";
            if (thinkStartIndex != -1) {
                thinkingProcess = rawResponse.substring(thinkStartIndex + 7, thinkEndIndex).trim();
            }
            
            // 提取实际答案（</think>之后的内容）
            String actualAnswer = rawResponse.substring(thinkEndIndex + 8).trim();
            
            return new ProcessedResponse(actualAnswer, thinkingProcess, rawResponse);
            
        } catch (Exception e) {
            // 解析出错，返回原始响应
            return new ProcessedResponse(rawResponse, "", rawResponse);
        }
    }
    
    /**
     * 处理后的响应对象
     */
    public static class ProcessedResponse {
        private final String answer;           // 实际答案
        private final String thinkingProcess;  // 思考过程
        private final String rawResponse;      // 原始响应
        
        public ProcessedResponse(String answer, String thinkingProcess, String rawResponse) {
            this.answer = answer;
            this.thinkingProcess = thinkingProcess;
            this.rawResponse = rawResponse;
        }
        
        public String getAnswer() {
            return answer;
        }
        
        public String getThinkingProcess() {
            return thinkingProcess;
        }
        
        public String getRawResponse() {
            return rawResponse;
        }
        
        public boolean hasThinkingProcess() {
            return thinkingProcess != null && !thinkingProcess.trim().isEmpty();
        }
        
        public String getDisplayText() {
            return answer != null && !answer.trim().isEmpty() ? answer : rawResponse;
        }
        
        @Override
        public String toString() {
            return "ProcessedResponse{" +
                    "answer='" + truncate(answer, 50) + '\'' +
                    ", hasThinking=" + hasThinkingProcess() +
                    ", rawLength=" + (rawResponse != null ? rawResponse.length() : 0) +
                    '}';
        }
        
        private String truncate(String str, int maxLength) {
            if (str == null) return "null";
            if (str.length() <= maxLength) return str;
            return str.substring(0, maxLength) + "...";
        }
    }
    
    /**
     * 示例用法
     */
    public static void main(String[] args) {
        // 测试不同类型的响应
        
        // 1. 包含思考过程的响应
        String thinkingResponse = "<think>\n这是思考过程\n用户问1+1等于几\n答案是2\n</think>\n\n1+1等于2。";
        ProcessedResponse processed1 = processResponse(thinkingResponse);
        System.out.println("测试1 - 包含思考过程:");
        System.out.println("实际答案: " + processed1.getAnswer());
        System.out.println("思考过程: " + processed1.getThinkingProcess());
        System.out.println("有思考过程: " + processed1.hasThinkingProcess());
        System.out.println();
        
        // 2. 普通响应
        String normalResponse = "Java是一种编程语言。";
        ProcessedResponse processed2 = processResponse(normalResponse);
        System.out.println("测试2 - 普通响应:");
        System.out.println("实际答案: " + processed2.getAnswer());
        System.out.println("有思考过程: " + processed2.hasThinkingProcess());
        System.out.println();
        
        // 3. 空响应
        ProcessedResponse processed3 = processResponse("");
        System.out.println("测试3 - 空响应:");
        System.out.println("实际答案: '" + processed3.getAnswer() + "'");
        System.out.println("有思考过程: " + processed3.hasThinkingProcess());
    }
}
